#include "Graphics/sdlgpu.h"
#include "Graphics/renderer.h"
#include "SDL3_image/SDL_image.h"
#include "glm/glm.hpp"
#include "glm/gtc/matrix_transform.hpp"

using namespace Jin;

struct Sprite {
	glm::vec2 pos;
	glm::vec2 size;
	glm::vec4 rotation;
	glm::vec4 texcoord;
	glm::vec4 color;
};

struct SpriteVertex {
	glm::vec4 position;
	glm::vec4 texcoord;
	glm::vec4 color;
};

struct AppContext final {

	enum {
		WIN_WIDTH = 800,
		WIN_HEIGHT = 600,
		WIN_FLAGS = 0,

		SPRITE_COUNT = 512
	};

	SDLFramework sdlfx;
	SDLWindowUPtr mainWindow;
	GPUDeviceUPtr gpuDevice;
	GPUGraphicsPipelineUPtr graphicsPipeline;
	GPUComputePipelineUPtr computePipeline;
	GPUSamplerUPtr texSampler;
	GPUUploadBufferUPtr spritesUpload;
	GPUComputeReadBufferUPtr spritesBuffer;
	GPUComputeWriteVertexBufferUPtr spritesVertexBuffer;
	GPUIndexBufferUPtr spritesIndexBuffer;
	GPUTextureSPtr spritesTexture;
	glm::mat4 camera{ 1.0f };
	std::vector<Sprite> sprites;

	AppContext();
	~AppContext();
	bool Init();
	void Run();
	void Cleanup();

private:
	bool _CreateMainWindow();
	bool _CreateGPUDevice();
	bool _CreateGraphicsPipeline(stringv vertFilename, stringv fragFilename);
	bool _CreateComputePipeline(stringv compFilename);
	bool _CreateSampler();
	bool _CreateImageTexture(stringv filename);
	void _UpdateSprites();
};

AppContext::AppContext() {
	SDL_LogInfo(0, "SDL version: %d.%d.%d", sdlfx.version.major, sdlfx.version.minor, sdlfx.version.micor);
}

AppContext::~AppContext() {

}

bool AppContext::_CreateMainWindow() {
	mainWindow = std::make_unique<SDLWindow>("testSDLGPU", WIN_WIDTH, WIN_HEIGHT, WIN_FLAGS);
	if (!*mainWindow) {
		SDL_LogError(0, "Create window failed! error: %s", SDL_GetError());
		return false;
	}
	SDL_LogInfo(0, "Create window successed!");
	sdlfx.SetFPS(200);
	sdlfx.onFPS = [this](Uint32 fps) {
		mainWindow->SetTitle(std::format("fps:{}, delta:{}", fps, sdlfx.deltaTime).c_str());
		};
	return true;
}

bool AppContext::_CreateGPUDevice() {
	gpuDevice = std::make_unique<GPUDevice>(*mainWindow);
	if (!*gpuDevice) {
		SDL_LogError(0, "Create GPU Device failed! error: %s", SDL_GetError());
		return false;
	}
	SDL_LogInfo(0, "Create GPUDevice dirver: %s successed!", gpuDevice->GetDeviceDriver().data());
	gpuDevice->SetSwapchainParameters(SDL_GPU_SWAPCHAINCOMPOSITION_SDR, SDL_GPU_PRESENTMODE_IMMEDIATE);
	return true;
}

bool AppContext::_CreateGraphicsPipeline(stringv vertFilename, stringv fragFilename) {
	GPUVertexShaderInfo vertexShaderInfo;
	vertexShaderInfo.vertexBufferDescriptions.resize(1);
	vertexShaderInfo.vertexBufferDescriptions[0].pitch = sizeof(SpriteVertex);
	vertexShaderInfo.vertexAttributes.resize(3);
	vertexShaderInfo.vertexAttributes[0].location = 0;
	vertexShaderInfo.vertexAttributes[0].format = SDL_GPU_VERTEXELEMENTFORMAT_FLOAT4;
	vertexShaderInfo.vertexAttributes[0].offset = 0;
	vertexShaderInfo.vertexAttributes[1].location = 1;
	vertexShaderInfo.vertexAttributes[1].format = SDL_GPU_VERTEXELEMENTFORMAT_FLOAT4;
	vertexShaderInfo.vertexAttributes[1].offset = offsetof(SpriteVertex, texcoord);
	vertexShaderInfo.vertexAttributes[2].location = 2;
	vertexShaderInfo.vertexAttributes[2].format = SDL_GPU_VERTEXELEMENTFORMAT_FLOAT4;
	vertexShaderInfo.vertexAttributes[2].offset = offsetof(SpriteVertex, color);
	GPUFragmentShaderInfo fragmentShaderInfo;
	fragmentShaderInfo.numSamplers = 1;

	SDLFileStorage storage("shaders");
	if (!storage) {
		SDL_LogError(0, "open shaders storage failed! %s", SDL_GetError());
		return false;
	}
	if (!storage.ReadFile(vertFilename, vertexShaderInfo.code)) {
		SDL_LogError(0, "Read file: %s failed! %s", vertFilename.data(), SDL_GetError());
		return false;
	}
	if (!storage.ReadFile(fragFilename, fragmentShaderInfo.code)) {
		SDL_LogError(0, "Read file: %s failed! %s", fragFilename.data(), SDL_GetError());
		return false;
	}

	graphicsPipeline = std::make_unique<GPUGraphicsPipeline>(*gpuDevice, vertexShaderInfo, fragmentShaderInfo);
	if (!*graphicsPipeline) {
		SDL_LogError(0, "Creaet GraphicsPipeline failed! error: %s", SDL_GetError());
		return false;
	}
	SDL_LogInfo(0, "Create GraphicsPipeline successed!");
	return true;
}

bool AppContext::_CreateComputePipeline(stringv compFilename) {
	GPUComputeShaderInfo computeShaderInfo;
	computeShaderInfo.numUniformBuffers = 1;
	computeShaderInfo.numReadonlyStorageBuffers = 1;
	computeShaderInfo.numReadwriteStorageBuffers = 1;

	SDLFileStorage storage("shaders");
	if (!storage) {
		SDL_LogError(0, "open shaders storage failed! %s", SDL_GetError());
		return false;
	}
	if (!storage.ReadFile(compFilename, computeShaderInfo.code)) {
		SDL_LogError(0, "Read file: %s failed! %s", compFilename.data(), SDL_GetError());
		return false;
	}

	computePipeline = std::make_unique<GPUComputePipeline>(*gpuDevice, computeShaderInfo);
	if (!*computePipeline) {
		SDL_LogError(0, "Creaet ComputePipeline failed! error: %s", SDL_GetError());
		return false;
	}
	SDL_LogInfo(0, "Create ComputePipeline successed!");
	return true;
}

bool AppContext::_CreateSampler() {
	GPUSamplerCreateInfo createInfo{};
	createInfo.min_filter = SDL_GPU_FILTER_NEAREST;
	createInfo.mag_filter = SDL_GPU_FILTER_NEAREST;
	createInfo.mipmap_mode = SDL_GPU_SAMPLERMIPMAPMODE_NEAREST;
	texSampler = std::make_unique<GPUSampler>(*gpuDevice, createInfo);
	if (!*texSampler) {
		SDL_LogError(0, "Create Sampler failed! error: %s", SDL_GetError());
		return false;
	}
	SDL_LogInfo(0, "Create Sampler successed!");
	return true;
}

bool AppContext::_CreateImageTexture(stringv filename) {
	spritesTexture = std::make_shared<GPUTexture>(*gpuDevice, "images/" + string{ filename });
	if (!*spritesTexture) {
		SDL_LogError(0, "Create Image Texture failed! error: %s", SDL_GetError());
		return false;
	}
	SDL_LogInfo(0, "Create Image Texture successed!");
	return true;
}

void AppContext::_UpdateSprites() {
	static const glm::vec2 texcoords[4]{
	{ 0.0f, 0.0f },
	{ 0.5f, 0.0f },
	{ 0.0f, 0.5f },
	{ 0.5f, 0.5f }
	};

	sprites.resize(SPRITE_COUNT);
	Sprite* sprite = sprites.data();
	for (Uint32 i = 0; i < SPRITE_COUNT; ++i) {
		Sint32 index = SDL_rand(4);
		sprite[i].pos = { (float)(SDL_rand(WIN_WIDTH - 200) + 100), (float)(SDL_rand(WIN_HEIGHT - 200) + 100) };
		sprite[i].size = { 32, 32 };
		sprite[i].rotation.s = SDL_randf();
		sprite[i].texcoord = { texcoords[index].x, texcoords[index].y, 0.5f, 0.5f };
		sprite[i].color = { 1.0f, 1.0f, 1.0f, 1.0f };
	}
}

bool AppContext::Init() {
	if (!_CreateMainWindow()) {
		return false;
	}
	if (!_CreateGPUDevice()) {
		return false;
	}
	if (!_CreateGraphicsPipeline("shader7.vert.spv", "shader7.frag.spv")) {
		return false;
	}
	if (!_CreateComputePipeline("shader7.comp.spv")) {
		return false;
	}
	if (!_CreateSampler()) {
		return false;
	}
	if (!_CreateImageTexture("ravioli_atlas.bmp")) {
		return false;
	}

	spritesUpload = std::make_unique<GPUUploadBuffer>(*gpuDevice, Uint32(SPRITE_COUNT * sizeof(Sprite)));
	if (!*spritesUpload) {
		return false;
	}
	spritesBuffer = std::make_unique<GPUComputeReadBuffer>(*gpuDevice, Uint32(SPRITE_COUNT * sizeof(Sprite)));
	if (!*spritesBuffer) {
		return false;
	}
	else {
		_UpdateSprites();
	}

	spritesVertexBuffer = std::make_unique<GPUComputeWriteVertexBuffer>(*gpuDevice, Uint32(SPRITE_COUNT * 4 * sizeof(SpriteVertex)));
	if (!*spritesVertexBuffer) {
		return false;
	}
	spritesIndexBuffer = std::make_unique<GPUIndexBuffer>(*gpuDevice, Uint32(SPRITE_COUNT * 6 * sizeof(Uint32)));
	if (!*spritesIndexBuffer) {
		return false;
	}
	else {
		GPUCommandBuffer cmd{ *gpuDevice };
		if (!cmd) {
			return false;
		}
		GPUUploadBuffer uploadBuffer{ *gpuDevice, Uint32(SPRITE_COUNT * 6 * sizeof(Uint32)) };
		GPUTransferBufferMapperT<Uint32*> mapper{ uploadBuffer };
		Uint32* index = mapper;
		for (Uint32 i = 0, j = 0; i < SPRITE_COUNT * 6; i += 6, j += 4) {
			index[i] = j;
			index[i + 1] = j + 1;
			index[i + 2] = j + 2;
			index[i + 3] = j + 3;
			index[i + 4] = j + 2;
			index[i + 5] = j + 1;
		}
		GPUCopyPass{ cmd }.UploadToGPUBuffer({ uploadBuffer, 0 }, { *spritesIndexBuffer, 0, SPRITE_COUNT * 6 * sizeof(Uint32) });
	}

	SDL_srand(0);

	camera = glm::ortho(0.0f, (float)WIN_WIDTH, (float)WIN_HEIGHT, 0.0f, 0.0f, -1.0f);

	return true;
}

void AppContext::Run() {
	static const glm::vec2 texcoords[4]{
		{ 0.0f, 0.0f },
		{ 0.5f, 0.0f },
		{ 0.0f, 0.5f },
		{ 0.5f, 0.5f }
	};

	sdlfx.onFrame = [this]() {
		if (mainWindow->GetWindowFlags() & SDL_WINDOW_MINIMIZED) {
			return;
		}
		GPUCommandBuffer cmd{ *gpuDevice };
		if (!cmd) {
			return;
		}
		GPUSwapchainTexture swapchainTexture{ cmd, *mainWindow };
		if (!swapchainTexture) {
			return;
		}

		cmd.PushComputeUniformData(0, &camera, sizeof(camera));
		
		{
			static float angle = 0;
			angle += sdlfx.deltaTime;
			Sprite* sprite = sprites.data();
			for (Uint32 i = 0; i < SPRITE_COUNT; ++i) {
				sprite[i].rotation.s = angle;
			}
			GPUTransferBufferMapperT<Sprite*> mapper{ *spritesUpload, true };
			memcpy(mapper, sprites.data(), SPRITE_COUNT * sizeof(Sprite));
			GPUCopyPass{ cmd }.UploadToGPUBuffer({ *spritesUpload, 0 }, { *spritesBuffer, 0, SPRITE_COUNT * sizeof(Sprite) }, true);

			GPUComputePass computePass{ cmd, GPUStorageBufferReadWriteBinding{ *spritesVertexBuffer, true } };
			computePass.BindComputePipeline(*computePipeline);
			computePass.BindComputeStorageBuffers(0, *spritesBuffer);
			computePass.DispatchGPUCompute((Uint32)SDL_ceilf((float)SPRITE_COUNT / 64.0f), 1, 1);
		}

		GPUColorTargetInfo colorTargetInfo{};
		colorTargetInfo.texture = swapchainTexture;
		colorTargetInfo.clear_color = { 0.2f, 0.3f, 0.3f, 1.0f };
		colorTargetInfo.load_op = SDL_GPU_LOADOP_CLEAR;
		colorTargetInfo.store_op = SDL_GPU_STOREOP_STORE;
		GPURenderPass renderPass{ cmd, &colorTargetInfo, 1 };
		if (!renderPass) {
			return;
		}
		renderPass.BindGraphicsPipeline(*graphicsPipeline);
		renderPass.BindVertexBuffers(0, { *spritesVertexBuffer });
		renderPass.BindIndexBuffer({ *spritesIndexBuffer }, SDL_GPU_INDEXELEMENTSIZE_32BIT);
		renderPass.BindFragmentSampler(0, { *spritesTexture, *texSampler });
		renderPass.DrawIndexedPrimitives(SPRITE_COUNT * 6, 1, 0, 0, 0);
		};
	sdlfx.onEvent = [this](EventData& event) {
		if (event.type == SDL_EVENT_KEY_DOWN) {
			if (event.key.key == SDLK_G) {
				_UpdateSprites();
			}
		}
		return true;
		};
	sdlfx.Run();
}


void AppContext::Cleanup() {
	gpuDevice->WaitIdle();
}

int main() {
	try {
		AppContext app;
		if (app.Init()) {
			app.Run();
		}
		app.Cleanup();
	}
	catch (const std::exception& e) {
		SDL_LogError(0, "App error! '%s'", e.what());
		return -1;
	}
	return 0;
}
